from torch.utils.data.dataset import Dataset
from PIL import Image


def convert_input_shape(input_shape):
    return input_shape if isinstance(input_shape, tuple) else (input_shape, input_shape)


def class_dict(class_names):
    label_dict = {}
    for i, word in enumerate(class_names):
        label_dict[word] = i
    return label_dict


class ViTDataSet(Dataset):
    def __init__(self, data_list, input_shape, class_names, transform=None):
        super(ViTDataSet, self).__init__()
        self.length = len(data_list)
        self.data_list = data_list
        self.input_shape = convert_input_shape(input_shape)
        self.classes_names = class_dict(class_names)
        self.transform = transform

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        img_path = self.data_list[idx]
        img = Image.open(img_path)
        img_transformed = self.transform(img)
        image_label = img_path.split("\\")[-1].split(".")[0]
        label = self.classes_names[image_label]

        return img_transformed, label
